from nnmnkwii.datasets import FileDataSource, FileSourceDataset
from nnmnkwii.preprocessing import minmax, meanvar, minmax_scale, scale
from os.path import join
from glob import glob
import numpy as np
from sklearn.model_selection import train_test_split
import config as cfg
from functools import partial


class BinaryFileSource(FileDataSource):
    def __init__(self, data_root, dim, train):
        self.data_root = data_root
        self.dim = dim
        self.train = train

    def collect_files(self):
        files = sorted(glob(join(self.data_root, "*.bin")))
        # print('len files', len(files))
        files = files[:len(files) - 5]  # last 5 is real testset
        train_files, test_files = train_test_split(files,
                                                   test_size=cfg.test_size,
                                                   random_state=cfg.random_state)
        if self.train:
            return train_files
        else:
            return test_files

    def collect_features(self, path):
        return np.fromfile(path, dtype=np.float32).reshape(-1, self.dim)


class LoadedData(object):

    def __init__(self):
        self.x = None
        self.y = None
        self.utt_lengths = None
        self.x_min = None
        self.x_max = None
        self.y_mean = None
        self.y_var = None
        self.y_scale = None

    def load(self):
        x = {"duration": {}, "acoustic": {}}
        y = {"duration": {}, "acoustic": {}}
        utt_lengths = {"duration": {}, "acoustic": {}}
        for ty in ["duration", "acoustic"]:
            for phase in ["train", "test"]:
                train = phase == "train"
                x_dim = cfg.duration_linguistic_dim if ty == "duration" else cfg.acoustic_linguisic_dim
                y_dim = cfg.duration_dim if ty == "duration" else cfg.acoustic_dim
                x[ty][phase] = list(FileSourceDataset(
                    BinaryFileSource(join(cfg.DATA_ROOT, "X_{}".format(ty)),
                                     dim=x_dim,
                                     train=train)))
                y[ty][phase] = list(FileSourceDataset(
                    BinaryFileSource(join(cfg.DATA_ROOT, "Y_{}".format(ty)),
                                     dim=y_dim,
                                     train=train)))
                # this triggers file loads, but can be neglectable in terms of performance.
                utt_lengths[ty][phase] = [len(x) for x in x[ty][phase]]

        x_min = {}
        x_max = {}
        y_mean = {}
        y_var = {}
        y_scale = {}

        for typ in ["acoustic", "duration"]:
            x_min[typ], x_max[typ] = minmax(x[typ]["train"],
                                            utt_lengths[typ]["train"])
            y_mean[typ], y_var[typ] = meanvar(y[typ]["train"],
                                              utt_lengths[typ]["train"])
            y_scale[typ] = np.sqrt(y_var[typ])

        self.x = x
        self.y = y
        self.utt_lengths = utt_lengths
        self.x_min = x_min
        self.x_max = x_max
        self.y_mean = y_mean
        self.y_var = y_var
        self.y_scale = y_scale

    def norm(self):
        for typ in ["acoustic", "duration"]:
            x_norm_func = partial(minmax_scale, data_min=self.x_min[typ],
                                  data_max=self.x_max[typ],
                                  feature_range=(0.01, 0.99))
            y_norm_func = partial(scale, data_mean=self.y_mean[typ],
                                  data_std=self.y_scale[typ])
            for phase in ["train", "test"]:
                self.x[typ][phase] = [x_norm_func(ele) for ele in
                                      self.x[typ][phase]]
                self.y[typ][phase] = [y_norm_func(ele) for ele in
                                      self.y[typ][phase]]
